import numpy as np
from scipy.optimize import root_scalar


def find_speed_from_flow_rate(flow_rate, coefficients):
    """
    根据给定的出水量和多项式系数，计算对应的转速。

    参数:
    - flow_rate (float): 出水量
    - coefficients (list or array-like): 多项式系数，从最高次到最低次

    返回:
    - speed (float): 对应的转速，如果存在多个实数根，则返回最合理的正值
    """
    # 创建多项式函数
    poly_func = np.poly1d(coefficients)

    # 定义目标函数 f(x) = ax^3 + bx^2 + cx + d - flow_rate
    def target_function(speed):
        return poly_func(speed) - flow_rate

    # 估计初始猜测值范围
    lower_bound = 0  # 转速不能为负
    upper_bound = 600  # 假设转速不会超过这个值，可以根据实际情况调整

    try:
        # 使用 root_scalar 寻找根，提供上下界以提高稳定性
        result = root_scalar(target_function, bracket=[lower_bound, upper_bound], method='brentq')
        if result.converged:
            return result.root
        else:
            raise ValueError("Root finding did not converge.")
    except ValueError as e:
        print(f"Error in finding the root: {e}")
        return None


# 示例调用
if __name__ == "__main__":
    # 给定的多项式系数（从最高次到最低次）
    coefficients = [-2.97967366e-06, 1.70126107e-03, 1.00979709e+00, 4.01233333e+00]

    try:
        # 输入示例出水量
        flow_rate = float(input("请输入出水量 (L/min): "))
        speed = find_speed_from_flow_rate(flow_rate, coefficients)
        if speed is not None:
            print(f"对应的转速为: {speed:.2f} rpm")
        else:
            print("无法找到有效的转速。")
    except ValueError as e:
        print(e)